package cn.jly.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, LongType, StructField, StructType}

/**
 * @author lanyangji
 * @date 2019/11/30 22:28
 */
object Spark04_UDAF {

  def main(args: Array[String]): Unit = {

    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("Spark04_UDAF")
    val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()

    import spark.implicits._

    // 自定义聚合函数 -- 这边模拟average函数
    // 注册函数
    spark.udf.register("myAvg", new MyAvg)

    // 读取数据创建df
    val df: DataFrame = spark.read.json("input/people.json")

    // 建表
    df.createOrReplaceTempView("persons")
    df.show()

    // 使用自定义的udaf函数
    spark.sql("select myAvg(age) as avg_age from persons").show()

    // 释放资源
    spark.close()
  }
}

// 声明用户自定义的聚合函数--模拟average函数
class MyAvg extends UserDefinedAggregateFunction {

  /**
   * 输入的数据结构
   *
   * @return
   */
  override def inputSchema: StructType = {

    // 写法1
    //new StructType().add("age", LongType)
    // 写法2
    StructType(StructField("inputColumn", LongType) :: Nil)
  }

  /**
   * 计算时的数据结构
   * 弱类型，你不能弄错sum和count的顺序
   *
   * @return
   */
  override def bufferSchema: StructType = {

    // 写法1
    //new StructType().add("sum", LongType).add("count", LongType)
    // 写法2
    StructType(StructField("sum", LongType) :: StructField("count", LongType) :: Nil)
  }

  /**
   * 返回值的数据类型
   *
   * @return
   */
  override def dataType: DataType = DoubleType

  /**
   * 函数是否稳定，即相同的输入是不是一定返回相同的输出
   *
   * @return
   */
  override def deterministic: Boolean = true

  /**
   * 计算之前缓冲区的初始化
   *
   * @param buffer
   */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {

    // 不关注类型，只关注数据
    buffer(0) = 0L
    buffer(1) = 0L
  }

  /**
   * 相同executor间数据的合并
   *
   * @param buffer
   * @param input
   */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {

    // sum
    buffer(0) = buffer.getLong(0) + input.getLong(0)
    // count + 1
    buffer(1) = buffer.getLong(1) + 1
  }

  /**
   * 不同executor之间数据的合并
   *
   * @param buffer1
   * @param buffer2
   */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {

    // sum
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    // count
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  /**
   * 计算
   * res = sum/count
   *
   * @param buffer
   * @return
   */
  override def evaluate(buffer: Row): Any = buffer.getLong(0).toDouble / buffer.getLong(1).toDouble
}
